/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package thrift;

import org.apache.thrift.TProcessor;
import org.apache.thrift.transport.TNonblockingServerTransport;
import org.apache.thrift.transport.TNonblockingTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.spi.SelectorProvider;
import java.util.*;
import java.util.concurrent.*;

/**
 * Copyright (C) 2018
 * All rights reserved
 * User: yulong.zhang
 * Date:2018年11月23日11:13:33
 * @author Administrator
 */
public class KoalasThreadedSelectorServer extends KoalasAbstractNonblockingServer {
    private static final Logger LOGGER = LoggerFactory.getLogger(KoalasThreadedSelectorServer.class.getName());

    public static class Args extends AbstractNonblockingServerArgs<Args> {

        public int selectorThreads = 2;

        private int workerThreads = 5;
        private int stopTimeoutVal = 60;
        private TimeUnit stopTimeoutUnit = TimeUnit.SECONDS;
        private ExecutorService executorService = null;

        private int acceptQueueSizePerThread = 4;

        public static enum AcceptPolicy {

            FAIR_ACCEPT,

            FAST_ACCEPT
        }

        private AcceptPolicy acceptPolicy = AcceptPolicy.FAST_ACCEPT;

        public Args(TNonblockingServerTransport transport) {
            super(transport);
        }

        public Args selectorThreads(int i) {
            selectorThreads = i;
            return this;
        }

        public int getSelectorThreads() {
            return selectorThreads;
        }

        public Args workerThreads(int i) {
            workerThreads = i;
            return this;
        }

        public int getWorkerThreads() {
            return workerThreads;
        }

        public int getStopTimeoutVal() {
            return stopTimeoutVal;
        }

        public Args stopTimeoutVal(int stopTimeoutVal) {
            this.stopTimeoutVal = stopTimeoutVal;
            return this;
        }

        public TimeUnit getStopTimeoutUnit() {
            return stopTimeoutUnit;
        }

        public Args stopTimeoutUnit(TimeUnit stopTimeoutUnit) {
            this.stopTimeoutUnit = stopTimeoutUnit;
            return this;
        }

        public ExecutorService getExecutorService() {
            return executorService;
        }

        public Args executorService(ExecutorService executorService) {
            this.executorService = executorService;
            return this;
        }

        public int getAcceptQueueSizePerThread() {
            return acceptQueueSizePerThread;
        }

        public Args acceptQueueSizePerThread(int acceptQueueSizePerThread) {
            this.acceptQueueSizePerThread = acceptQueueSizePerThread;
            return this;
        }

        public AcceptPolicy getAcceptPolicy() {
            return acceptPolicy;
        }

        public Args acceptPolicy(AcceptPolicy acceptPolicy) {
            this.acceptPolicy = acceptPolicy;
            return this;
        }

        public void validate() {
            if (selectorThreads <= 0) {
                throw new IllegalArgumentException("selectorThreads must be positive.");
            }
            if (workerThreads < 0) {
                throw new IllegalArgumentException("workerThreads must be non-negative.");
            }
            if (acceptQueueSizePerThread <= 0) {
                throw new IllegalArgumentException("acceptQueueSizePerThread must be positive.");
            }
        }
    }

    private volatile boolean stopped_ = true;

    /**
     * 连接线程
     */
    private AcceptThread acceptThread;

    /**
     * 读写线程
     */
    private final Set<SelectorThread> selectorThreads = new HashSet<SelectorThread>();

    private final ExecutorService invoker;

    private final Args args;

    private String privateKey;
    private String publicKey;
    private String serviceName;
    private TProcessor tGenericProcessor;

    private boolean cat;

    public boolean isCat() {
        return cat;
    }

    public void setCat(boolean cat) {
        this.cat = cat;
    }

    public TProcessor gettGenericProcessor() {
        return tGenericProcessor;
    }

    public void settGenericProcessor(TProcessor tGenericProcessor) {
        this.tGenericProcessor = tGenericProcessor;
    }

    public String getServiceName() {
        return serviceName;
    }

    public void setServiceName(String serviceName) {
        this.serviceName = serviceName;
    }

    public String getPrivateKey() {
        return privateKey;
    }

    public void setPrivateKey(String privateKey) {
        this.privateKey = privateKey;
    }

    public String getPublicKey() {
        return publicKey;
    }

    public void setPublicKey(String publicKey) {
        this.publicKey = publicKey;
    }

    public KoalasThreadedSelectorServer(Args args) {
        super(args);
        args.validate();
        invoker = args.executorService == null ? createDefaultExecutor(args) : args.executorService;
        this.args = args;
    }

    @Override
    protected boolean startThreads() {
        try {
            for (int i = 0; i < args.selectorThreads; ++i) {
                selectorThreads.add(new SelectorThread(args.acceptQueueSizePerThread));
            }
            acceptThread = new AcceptThread((TNonblockingServerTransport) serverTransport_,
                    createSelectorThreadLoadBalancer(selectorThreads));
            stopped_ = false;
            for (SelectorThread thread : selectorThreads) {
                thread.start();
            }
            acceptThread.start();
            return true;
        } catch (IOException e) {
            LOGGER.error("Failed to start threads!", e);
            return false;
        }
    }

    @Override
    protected void waitForShutdown() {
        try {
            joinThreads();
        } catch (InterruptedException e) {
            LOGGER.error("Interrupted while joining threads!", e);
        }
        gracefullyShutdownInvokerPool();
    }

    protected void joinThreads() throws InterruptedException {
        acceptThread.join();
        for (SelectorThread thread : selectorThreads) {
            thread.join();
        }
    }

    @Override
    public void stop() {
        stopped_ = true;
        stopListening();

        if (acceptThread != null) {
            acceptThread.wakeupSelector();
        }
        if (selectorThreads != null) {
            for (SelectorThread thread : selectorThreads) {
                if (thread != null) {
                    thread.wakeupSelector();
                }
            }
        }
    }

    protected void gracefullyShutdownInvokerPool() {
        invoker.shutdown();

        long timeoutMS = args.stopTimeoutUnit.toMillis(args.stopTimeoutVal);
        long now = System.currentTimeMillis();
        while (timeoutMS >= 0) {
            try {
                invoker.awaitTermination(timeoutMS, TimeUnit.MILLISECONDS);
                break;
            } catch (InterruptedException ix) {
                long newnow = System.currentTimeMillis();
                timeoutMS -= (newnow - now);
                now = newnow;
            }
        }
    }

    /**
     * 请求代理执行，这里和thrift的TNonblockingServer不同
     * TNonblockingServer由单线程执行frameBuffer.invoke();
     * 而这里改写之后则是交由线程池去执行frameBuffer.invoke();
     * @param frameBuffer
     * @return
     */
    @Override
    protected boolean requestInvoke(FrameBuffer frameBuffer) {
        Runnable invocation = getRunnable(frameBuffer);
        if (invoker != null) {
            try {
                invoker.execute(invocation);
                return true;
            } catch (RejectedExecutionException rx) {
                LOGGER.warn("ExecutorService rejected execution!", rx);
                return false;
            }
        } else {
            invocation.run();
            return true;
        }
    }

    protected Runnable getRunnable(FrameBuffer frameBuffer) {
        return new Invocation(frameBuffer);
    }

    /**
     * 创建默认线程池
     * @param options
     * @return
     */
    protected static ExecutorService createDefaultExecutor(Args options) {
        return (options.workerThreads > 0) ? Executors.newFixedThreadPool(options.workerThreads) : null;
    }

    private static BlockingQueue<TNonblockingTransport> createDefaultAcceptQueue(int queueSize) {
        if (queueSize == 0) {
            return new LinkedBlockingQueue<TNonblockingTransport>();
        }
        return new ArrayBlockingQueue<TNonblockingTransport>(queueSize);
    }

    /**
     * 连接线程
     */
    protected class AcceptThread extends Thread {

        /**
         * TTransport主要处理服务器端和客户端的网络读写
         * TNonblockingServerTransport 以非阻塞方式操作的服务器传输。
         */
        private final TNonblockingServerTransport serverTransport;
        /**
         * 连接通道选择器
         */
        private final Selector acceptSelector;

        /**
         * 获取读写线程负载均衡器
         */
        private final SelectorThreadLoadBalancer threadChooser;

        /**
         * 构造方法赋值
         * @param serverTransport
         * @param threadChooser
         * @throws IOException
         */
        public AcceptThread(TNonblockingServerTransport serverTransport,
                            SelectorThreadLoadBalancer threadChooser) throws IOException {
            this.serverTransport = serverTransport;
            this.threadChooser = threadChooser;
            // 获取选择器
            this.acceptSelector = SelectorProvider.provider().openSelector();
            this.serverTransport.registerSelector(acceptSelector);
        }

        @Override
        public void run() {
            try {
                // 服务开启状态，负载获取读写线程
                while (!stopped_) {
                    select();
                }
            } catch (Throwable t) {
                LOGGER.error("run() exiting due to uncaught error", t);
            } finally {
                KoalasThreadedSelectorServer.this.stop();
            }
        }

        public void wakeupSelector() {
            acceptSelector.wakeup();
        }

        private void select() {
            try {
                // 选择一些I/O操作已经准备好的管道。每个管道对应着一个key。这个方法是一个阻塞的选择操作。
                // 当至少有一个通道被选择时才返回。当这个方法被执行时，当前线程是允许被中断的。
                acceptSelector.select();

                // 获取通道key集合
                Iterator<SelectionKey> selectedKeys = acceptSelector.selectedKeys().iterator();
                // 服务开启转态并且有准备好的通道
                while (!stopped_ && selectedKeys.hasNext()) {
                    SelectionKey key = selectedKeys.next();
                    /**
                     * 这里是处理连接线程，这里处理了这个key后，将其删除，避免下次监听依然监听到该链接
                     */
                    selectedKeys.remove();

                    if (!key.isValid()) {
                        continue;
                    }
                    // 有新的连接处理连接
                    if (key.isAcceptable()) {
                        handleAccept();
                    } else {
                        LOGGER.warn("Unexpected state in select! " + key.interestOps());
                    }
                }
            } catch (IOException e) {
                LOGGER.warn("Got an IOException while selecting!", e);
            }
        }

        private void handleAccept() {
            final TNonblockingTransport client = doAccept();
            if (client != null) {
                final SelectorThread targetThread = threadChooser.nextThread();

                if (args.acceptPolicy == Args.AcceptPolicy.FAST_ACCEPT || invoker == null) {
                    // 处理新增连接
                    doAddAccept(targetThread, client);
                } else {
                    try {
                        invoker.submit(new Runnable() {
                            @Override
                            public void run() {
                                doAddAccept(targetThread, client);
                            }
                        });
                    } catch (RejectedExecutionException rx) {
                        LOGGER.warn("ExecutorService rejected accept registration!", rx);
                        client.close();
                    }
                }
            }
        }

        /**
         * 获取thrift服务器异步数据传输器
         * @return
         */
        private TNonblockingTransport doAccept() {
            try {
                return (TNonblockingTransport) serverTransport.accept();
            } catch (TTransportException tte) {
                LOGGER.warn("Exception trying to accept!", tte);
                return null;
            }
        }

        private void doAddAccept(SelectorThread thread, TNonblockingTransport client) {
            // 保存到队列失败关闭当前连接
            if (!thread.addAcceptedConnection(client)) {
                client.close();
            }
        }
    }

    /**
     * 读写线程
     */
    protected class SelectorThread extends AbstractSelectThread {

        /**
         * 保存连接线程队列
         */
        private final BlockingQueue<TNonblockingTransport> acceptedQueue;

        public SelectorThread() throws IOException {
            this(new LinkedBlockingQueue<TNonblockingTransport>());
        }

        public SelectorThread(int maxPendingAccepts) throws IOException {
            this(createDefaultAcceptQueue(maxPendingAccepts));
        }

        public SelectorThread(BlockingQueue<TNonblockingTransport> acceptedQueue) throws IOException {
            this.acceptedQueue = acceptedQueue;
        }

        /**
         * 新增连接线程传输对象到队列
         * @param accepted
         * @return
         */
        public boolean addAcceptedConnection(TNonblockingTransport accepted) {
            try {
                acceptedQueue.put(accepted);
            } catch (InterruptedException e) {
                LOGGER.warn("Interrupted while adding accepted connection!", e);
                return false;
            }
            selector.wakeup();
            return true;
        }

        @Override
        public void run() {
            try {
                while (!stopped_) {
                    select();
                    // 处理接收连接（连接事件通道），注册可读时间监听
                    processAcceptedConnections();
                    processInterestChanges();
                }
                for (SelectionKey selectionKey : selector.keys()) {
                    cleanupSelectionKey(selectionKey);
                }
            } catch (Throwable t) {
                LOGGER.error("run() exiting due to uncaught error", t);
            } finally {
                // This will wake up the accept thread and the other selector threads
                KoalasThreadedSelectorServer.this.stop();
            }
        }

        private void select() {
            try {
                selector.select();

                Iterator<SelectionKey> selectedKeys = selector.selectedKeys().iterator();
                while (!stopped_ && selectedKeys.hasNext()) {
                    SelectionKey key = selectedKeys.next();
                    selectedKeys.remove();

                    if (!key.isValid()) {
                        cleanupSelectionKey(key);
                        continue;
                    }

                    if (key.isReadable()) {
                        // 处理读通道
                        handleRead(key);
                    } else if (key.isWritable()) {
                        // 处理写通道
                        handleWrite(key);
                    } else {
                        LOGGER.warn("Unexpected state in select! " + key.interestOps());
                    }
                }
            } catch (IOException e) {
                LOGGER.warn("Got an IOException while selecting!", e);
            }
        }

        private void processAcceptedConnections() {
            // 注册已接受的连接
            while (!stopped_) {
                TNonblockingTransport accepted = acceptedQueue.poll();
                if (accepted == null) {
                    break;
                }
                registerAccepted(accepted);
            }
        }

        private void registerAccepted(TNonblockingTransport accepted) {
            SelectionKey clientKey = null;
            try {
                clientKey = accepted.registerSelector(selector, SelectionKey.OP_READ);

                FrameBuffer frameBuffer = new FrameBuffer(accepted, clientKey, SelectorThread.this, privateKey, publicKey, serviceName, tGenericProcessor, cat);
                clientKey.attach(frameBuffer);
            } catch (IOException e) {
                LOGGER.warn("Failed to register accepted connection to selector!", e);
                if (clientKey != null) {
                    cleanupSelectionKey(clientKey);
                }
                accepted.close();
            }
        }
    }

    /**
     * 创建读写线程负载均衡器
     * @param threads
     * @return
     */
    protected SelectorThreadLoadBalancer createSelectorThreadLoadBalancer(Collection<? extends SelectorThread> threads) {
        return new SelectorThreadLoadBalancer(threads);
    }

    /**
     * 读写线程负载均衡器
     */
    protected class SelectorThreadLoadBalancer {
        private final Collection<? extends SelectorThread> threads;
        private Iterator<? extends SelectorThread> nextThreadIterator;

        public <T extends SelectorThread> SelectorThreadLoadBalancer(Collection<T> threads) {
            if (threads.isEmpty()) {
                throw new IllegalArgumentException("At least one selector thread is required");
            }
            this.threads = Collections.unmodifiableList(new ArrayList<T>(threads));
            nextThreadIterator = this.threads.iterator();
        }

        public SelectorThread nextThread() {
            if (!nextThreadIterator.hasNext()) {
                nextThreadIterator = threads.iterator();
            }
            return nextThreadIterator.next();
        }
    }
}
